#include "sfm_tracks.h"
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/connected_components.hpp>
namespace ville3d {

TracksBuilder::TracksBuilder() {}

TracksBuilder::~TracksBuilder() {}

bool TracksBuilder::Build(
    const std::vector<std::pair<IndexT, IndexT>> image_pair,
    const std::vector<Matches> matches, TrackSet& tracks) {
  std::vector<std::pair<IndexT, IndexT>> pts;
  std::vector<std::pair<IndexT, IndexT>> mpts;

  for (std::size_t i = 0; i < image_pair.size(); i++) {
    for (std::size_t j = 0; j < matches[i].size(); j++) {
      pts.emplace_back(image_pair[i].first, IndexT(matches[i][j].queryIdx));
      mpts.emplace_back(image_pair[i].second, IndexT(matches[i][j].trainIdx));
    }
  }
  return Build(pts, mpts, tracks);
}

bool TracksBuilder::Build(const std::vector<std::pair<IndexT, IndexT>>& pts,
                          const std::vector<std::pair<IndexT, IndexT>>& mpts,
                          TrackSet& tracks) {
  if (pts.size() != mpts.size()) {
    return false;
  }

  // create node
  std::set<std::pair<IndexT, IndexT>> features;
  for (std::size_t i = 0; i < pts.size(); i++) {
    features.insert(pts[i]);
    features.insert(mpts[i]);
  }

  // create graph  vertexIndex and indexVertex
  using namespace boost;
  adjacency_list<vecS, vecS, undirectedS, std::size_t, std::size_t> g;

  std::size_t index = 0;
  std::map<std::pair<IndexT, IndexT>, std::size_t> vertexIndex;
  std::map<std::size_t, std::pair<IndexT, IndexT>> indexVertex;

  for (auto pos = features.begin(); pos != features.end(); ++pos) {
    vertexIndex[*pos] = index;
    indexVertex[index] = *pos;
    index++;
  }

  for (std::size_t i = 0; i < pts.size(); i++) {
    add_edge(vertexIndex[pts[i]], vertexIndex[mpts[i]], g);
  }

  std::vector<int> comp(num_vertices(g));
  int num = connected_components(g, comp.data());

  // create  Trackset
  for (std::size_t i = 0; i < comp.size(); i++) {
    IndexT track_id = IndexT(comp[i]);
    IndexT image_id = IndexT(indexVertex[i].first);
    IndexT feat_id = IndexT(indexVertex[i].second);

    // for {track_id, { image_id { feat} } }
    tracks.tracks[track_id].emplace_back(image_id, feat_id);
    tracks.image_tracks[image_id].insert(track_id);
    tracks.imageFeat_tracks[{image_id, feat_id}] = track_id;
  }
  return num > 0;
}

bool TracksBuilder::Filter(TrackSet& tracks, std::size_t nLengthSupTo /*= 2*/) {
  // {track_id, {image_id, image_id, ...}}
  std::map<IndexT, std::set<IndexT>> tracks_container;  

  std::set<TrackID> problematic_track_id;
  for (auto& imageFeat_track : tracks.imageFeat_tracks) {
    std::pair<IndexT, IndexT> image_feat = imageFeat_track.first;
    TrackID track_id = imageFeat_track.second;
    if (tracks_container[track_id].insert(image_feat.first).second == false) {
      problematic_track_id.insert(track_id);  // invalid
    }
  }

  // Reject tracks that have too few observations
  for (const auto& val : tracks_container) {
    if (val.second.size() < nLengthSupTo) {
      problematic_track_id.insert(val.first);
    }
  }

  for (auto& track_id : problematic_track_id) {
    tracks.remove_track(track_id);
  }

  return true;
}

}  // namespace ville3d